# -*- coding: utf-8 -*-

from WSModule import WSModule
from WSResult import WSResult
from WSOption import WSOption
from SpiderThread import SpiderThread
from SpiderJob import SpiderJob
from SpiderCommon import SpiderCommon
from Registry import Registry
from SpiderResult import SpiderResult

import os, datetime, urllib, hashlib, shutil, sqlite3, re

class SpiderException(Exception): pass

class Spider(WSModule):
	log_path = os.getcwd() + '/results/Spider-result'
	options = {"threads": WSOption("threads", "count thread", 5, False, ['-t', '--threads']),
			   "ignore": WSOption("ignore", "ignore regexp", "", False, ['-i', '--ignore-regexp']),
			   "only_one": WSOption("only_one", "only one", "", False, ['-o', '--only-one']),
			   "map": WSOption("map", "map name", "", False, ['-m', '--map-name']),
			   "url": WSOption("url", "target url", "", True, ['-u', '--url']),
			   "full-new": WSOption("full-new", "full new scan", "", False, ['-f', '--full-new']),
			   }

	def _prepare_main_page(self, site):
		coll = Registry.get('mongo').spider_urls
		coll.drop()

		data = {'hash': hashlib.md5(str(site.path + site.query).encode()).hexdigest(),
				'path': site.path,
				'query': site.query,
				'time': 0,
				'code':0,
				'checked': 0,
				'getted' : 0,
				'referer': '',
				'founder': 'spider'}

		coll.insert(data)

		coll.create_index([('hash', 1)], unique=True, dropDups=True)
		coll.create_index([('checked', 1)])

	def _options_to_registry(self, site):
		if self.options['ignore'].value:
			Registry.set('ignore_regexp', re.compile(self.options['ignore'].value))

		expr = ''
		for ext in Registry.get('config')['spider']['allow_links'].split(','):
			expr += '\.' + ext.strip() + '$|'
		expr = expr.rstrip('|')
		Registry.set('allow_regexp', re.compile(expr, re.I))

		if self.options['only_one'].value:
			tmp = self.options['only_one'].value.split("||")
			if len(tmp):
				only_one = []
				for regex in tmp:
					only_one.append({'regex': regex, 'block': False})
				Registry.set('only_one', only_one)

		Registry.set('map', site if not self.options['map'].value else self.options['map'].value)

	def run(self):
		self.result = SpiderResult()
		site = urllib.parse.urlparse(self.options['url'].value)
		self._options_to_registry(site.netloc)

		if self.options['full-new'].value:
			SpiderCommon.make_full_new_scan()

		if not os.path.exists(Registry.get('data_path') + Registry.get('map')):
			os.mkdir(Registry.get('data_path') + Registry.get('map'))
			os.chmod(Registry.get('data_path') + Registry.get('map'), 0o777)

		self._prepare_main_page(site)

		job = SpiderJob()

		workers = []
		for i in range(int(self.options['threads'].value)):
			workers.append(SpiderThread(job, site.netloc))
		self.kernel.createThreads(workers)

		while not self.kernel.finished():
			pass

		print("Total links count:" + str(Registry.get('mongo').spider_urls.count()))

	def help(self):
		return "Module Spider ver 1.0"
	def finished(self):
		return self.kernel.finished()

